import spacy
from collections import defaultdict
from config import Config

# 加载spacy中文模型（需要先下载：python -m spacy download zh_core_web_md）
try:
    nlp = spacy.load("zh_core_web_md")
except OSError:
    print("请先下载中文模型: python -m spacy download zh_core_web_md")
    nlp = None

# 预定义的实体词典（实际使用时需要扩展）
BRAND_DICT = {
    '华为': ['华为', 'Huawei', '问界', 'HarmonyOS', '鸿蒙'],
    'XM': ['小米', 'XM', 'Xiaomi', 'SU7', '雷军'],
    'LP': ['理想', 'Li Auto', '理想汽车', 'L7', 'L8', 'L9'],
    'XP': ['小鹏', 'XPeng', '小鹏汽车', 'G6', 'G9', 'P7'],
    'Geely': ['吉利', 'Geely', '极氪', 'Zeekr', '001', '007'],
    'LX': ['蔚来', 'NIO', 'Nio', 'ET7', 'ET5', 'ES6'],
    '博世': ['博世', 'Bosch'],
    '大陆': ['大陆集团', 'Continental'],
    '电装': ['电装', 'Denso'],
    '特斯拉': ['特斯拉', 'Tesla', 'Model Y', 'Model 3', '马斯克'],
    '丰田': ['丰田', 'Toyota', '凯美瑞', '卡罗拉'],
    '大众': ['大众', 'Volkswagen', 'VW', '奥迪', '奥迪', '保时捷'],
    '比亚迪': ['比亚迪', 'BYD', '仰望', '海豚', '海豹']
}

# PRODUCT_KEYWORDS = ['上市', '发布', '亮相', '预售', '降价', '召回', '改款', '换代']

def extract_entities(text):
    """使用spacy进行命名实体识别"""
    if not nlp or not text:
        return []

    doc = nlp(text)
    entities = []
    # 获取并打印NER管道组件的所有标签
    # 该代码行通过nlp对象获取名为"ner"的管道组件，然后访问其labels属性来获取所有可用的命名实体标签，
    # 最后将这些标签打印输出
    print(nlp.get_pipe("ner").labels)


    for ent in doc.ents:
        # 只关注组织、产品等实体类型
        if ent.label_ in ['ORG', 'PRODUCT', 'PERSON']:
            entities.append({
                'text': ent.text,
                'label': ent.label_,
                'start': ent.start_char,
                'end': ent.end_char
            })

    return entities


def map_entities_to_brands(entities):
    """将识别到的实体映射到预定义的品牌分类"""
    found_brands = set()
    found_products = set()

    for entity in entities:
        entity_text = entity['text']

        # 检查实体是否在预定义词典中
        for brand, keywords in BRAND_DICT.items():
            if any(keyword in entity_text for keyword in keywords):
                found_brands.add(brand)

        # 简单的产品识别规则（实际需要更复杂的逻辑）
        if '车' in entity_text or '版' in entity_text or '代' in entity_text:
            if len(entity_text) <= 10:  # 避免过长文本
                found_products.add(entity_text)

    # 优先级处理：优先返回产品
    if found_products:
        return list(found_products), 'product'
    elif found_brands:
        return list(found_brands), 'company'
    else:
        return [], 'unknown'